import keras
from keras import layers
from keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt
from AlexNet import AlexNet_v1

train_dir = './data/numbers/train'
validation_dir = './data/numbers/validation'

training_datagen = ImageDataGenerator(
    rescale=1./255,
    rotation_range=40,
    width_shift_range=0.2,
    height_shift_range=0.2,
    shear_range=0.2,
    horizontal_flip=True,
    fill_mode='nearest'
)

validation_datagen = ImageDataGenerator(
    rescale=1./255
)

training_generator = training_datagen.flow_from_directory(
    train_dir,
    target_size=(224,224),
    class_mode='categorical'
)

validation_generator = validation_datagen.flow_from_directory(
    validation_dir,
    target_size=(224,224),
    class_mode='categorical'
)

# model = keras.Sequential([
#     # Conv2D 卷积神经网络
#     # （3,3） 卷积核的大小
#
#     layers.Conv2D(64,(3,3),activation='relu',input_shape=(224,224,3),padding='SAME'),
#     layers.MaxPooling2D((2,2)),
#
#     layers.Conv2D(128,(3,3),activation='relu',padding='SAME',strides=1),
#     layers.MaxPooling2D((2,2)),
#
#     layers.Conv2D(128,(3,3),activation='relu',padding='SAME',strides=1),
#     layers.MaxPooling2D((2,2)),
#
#     layers.Flatten(),
#     layers.Dense(1000),
#
#     keras.layers.Dense(3,activation='softmax')
# ])

model = AlexNet_v1(224,224,3)

model.compile(
    loss=keras.losses.binary_crossentropy,
    optimizer="adam",
    metrics=['acc']
) #scipy
history = model.fit(training_generator,batch_size=20,epochs=100,validation_data=validation_generator,verbose=1)

# 保存模型到本地
model.save("test.keras")

acc = history.history['acc']
val_acc = history.history['val_acc']
loss = history.history['loss']
val_loss = history.history['val_loss']

epochs = range(len(acc))

plt.plot(epochs,acc,'r',label="Traing accuracy")
plt.plot(epochs,val_acc,'b',label="Validation accuracy")
plt.title('Training and Validation accuracy')
plt.legend(loc=0)
plt.figure()

plt.show()